#ifndef WEBSOCKET_H
#define WEBSOCKET_H

#include <string>
#include <functional>
#include <unordered_map>
#include <memory>
#include <websocketpp/config/asio_no_tls.hpp>
#include <websocketpp/server.hpp>

typedef websocketpp::server<websocketpp::config::asio> WebSocketServer;
typedef WebSocketServer::message_ptr MessagePtr;
typedef std::function<void(websocketpp::connection_hdl, const std::string&)> MessageHandler;

class WebSocketService {
public:
    WebSocketService(int port);
    ~WebSocketService();
    
    void run();
    void stop();
    void setMessageHandler(MessageHandler handler);
    void send(websocketpp::connection_hdl hdl, const std::string& message);
    void broadcast(const std::string& message);
    
    WebSocketServer& getServer() { return server_; }

private:
    void onOpen(websocketpp::connection_hdl hdl);
    void onClose(websocketpp::connection_hdl hdl);
    void onMessage(websocketpp::connection_hdl hdl, MessagePtr msg);
    
    WebSocketServer server_;
    int port_;
    MessageHandler messageHandler_;
    
    // 使用连接指针作为键
    std::unordered_map<void*, std::string> connections_;
};

#endif // WEBSOCKET_H